package babydisco.ICA;

import babydisco.Math.*;
import java.util.Arrays;

/**
 * 
 * @author rik_claessens
 */
public class FastICA {

    private static double[] meanValues;
    private static double[] eigenValues;
    private static double[][] zeroMeanVectors;
    private static double[][] covarianceMatrix;
    private static double[][] E;
    private static double[][] resultVectors;
    private static double[][] whiteningMatrix;
    private static double[][] dewhiteningMatrix;
    private static double[][] whitenedVectors;
    private static double[][] B;
    private static int maxIterations = 10;

    /**
     * Finds a certain number of indepentent components of the inout signal using fastica
     * @param input the input signals
     * @param maxIterations maximum number of iterations
     * @param epsilon level of accuracy
     * @param noComponents the number of components this method should return
     * @return the independent components of the signal
     */
    public static double[][] fastICA(double[][] input, int maxIterations, double epsilon, int noComponents) {
        doWhitening(input);
        int m = Matrix.getNumOfRows(whitenedVectors);
        int n = Matrix.getNumOfColumns(whitenedVectors);
        if (m > noComponents) {
            noComponents = m;
        }
        B = Matrix.random(noComponents, m);
        B = Matrix.mult(powerSymmMatrix(Matrix.square(B), -0.5), B);
        for (int k = 1; k < maxIterations; k++) {
            double[][] previousB = Matrix.clone(B);
            for (int c = 0; c < noComponents; c++) {
                //Step 1
                double[] previousW = Matrix.getVecOfRow(previousB, c);
                double[] eq1 = new double[m];

                for (int j = 0; j < n; j++) {                   
                    double one = Vector.dot(previousW, Matrix.getVecOfCol(whitenedVectors, j));
                    one = Math.pow(one, 3);
                    double[] two = Vector.scale(one, Matrix.getVecOfCol(whitenedVectors, j));
                    eq1 = Vector.add(eq1, two);
                }
                eq1 = Vector.scale((1.0 / (double) n), eq1);
                double[] eq2 = Vector.scale(-3, previousW);
                double[] w = Vector.add(eq1, eq2);              
                for (int j = 0; j < m; ++j) {
                    B[c][j] = w[j];
                }
            }
            B = Matrix.mult(powerSymmMatrix(Matrix.square(B), -0.5), B);
            double matrixDelta = deltaMatrices(B, previousB);
            if (matrixDelta < epsilon) {
                break;
            }
        }
        double[][] separatedMatrix = Matrix.mult(B, whiteningMatrix);
        return Matrix.mult(separatedMatrix, input);
    }

    /**
     * returns the mixingMatrix instead of the most significant components
     * @param input
     * @param maxIterations
     * @param epsilon
     * @param noComponents
     * @return
     */
    public static double[][] mixingMatrix(double[][] input, int maxIterations, double epsilon, int noComponents) {
        doWhitening(input);

        int m = Matrix.getNumOfRows(whitenedVectors);
        int n = Matrix.getNumOfColumns(whitenedVectors);

        if (m > noComponents) {
            noComponents = m;
        }


        B = Matrix.random(noComponents, m);
        B = Matrix.mult(powerSymmMatrix(Matrix.square(B), -0.5), B);

        for (int k = 1; k < maxIterations; k++) {               // Steps 2 - 4
            double[][] oldB = Matrix.clone(B);
            for (int c = 0; c < noComponents; c++) {
                //Step 1
                double[] prevW = Matrix.getVecOfRow(oldB, c);
                double[] eq1 = new double[m];

                for (int j = 0; j < n; j++) {

                    //First part of the equation
                    double one = Vector.dot(prevW, Matrix.getVecOfCol(whitenedVectors, j));
                    one = Math.pow(one, 3);

                    double[] two = Vector.scale(one, Matrix.getVecOfCol(whitenedVectors, j));
                    eq1 = Vector.add(eq1, two);
                }

                eq1 = Vector.scale((1.0 / (double) n), eq1);

                //Second part of the equation
                double[] eq2 = Vector.scale(-3, prevW);
                double[] w = Vector.add(eq1, eq2);
                //
                for (int j = 0; j < m; ++j) {
                    B[c][j] = w[j];
                }
            }


            // symmetric decorrelation by orthonormalisation
            B = Matrix.mult(
                    powerSymmMatrix(Matrix.square(B), -0.5),
                    B);

            double matrixDelta = deltaMatrices(B, oldB);
            System.out.println("Matrix delta: " + matrixDelta);
            if (matrixDelta < epsilon) {
                System.out.println("Converged after " + k + " iterations.");
                break;
            }
        }
        return B;
    }

    /**
     * Finds the major component of the input signal using fastica
     * 
     * @param input the input signals
     * @param maxIterations maximum number of iterations
     * @param epsilon level of accuracy
     * @return the major component
     */
    public static double[] fastICA(double[][] input, int maxIterations, double epsilon) {
        doWhitening(input);

        int m = Matrix.getNumOfRows(whitenedVectors);
        int n = Matrix.getNumOfColumns(whitenedVectors);

        double[] w = Vector.random(m);                              //Step 1
        w = Vector.normalize(w);

        for (int k = 1; k < maxIterations; k++) {               // Steps 2 - 4

            double[] prevW = Arrays.copyOf(w, w.length);

            double[] eq1 = new double[m];
            for (int j = 0; j < n; j++) {

                //First part of the equation
                double one = Vector.dot(prevW, Matrix.getVecOfCol(whitenedVectors, j));
                one = Math.pow(one, 3);
                double[] two = Vector.scale(one, Matrix.getVecOfCol(whitenedVectors, j));
                eq1 = Vector.add(eq1, two);
                //
            }

            eq1 = Vector.scale((1.0 / (double) n), eq1);

            //Second part of the equation
            double[] eq2 = Vector.scale(-3, prevW);
            //
            w = Vector.add(eq1, eq2);              // End of step 2
            w = Vector.normalize(w);                            // Step 3

            double dotTest = Math.abs(Vector.dot(w, prevW));
            if (dotTest >= (1 - epsilon)) {                     // Step 4
                break;
            }
        }

        double[][] wTransposed = new double[m][1];

        for (int i = 0; i < m; i++) {
            wTransposed[i][0] = w[i];
        }
        return Matrix.mult(wTransposed, whitenedVectors)[0];
    }

    public static void doWhitening(double[][] input) {
        // Centering, substract the mean from the signal vectors
        meanValues = calcMeanValues(input);

        zeroMeanVectors = Vector.addVecToSet(input, Vector.scale(-1.0, meanValues));

        // calculate the covariance matrix
        covarianceMatrix = Matrix.scale(Matrix.square(zeroMeanVectors), 1.0 / Matrix.getNumOfColumns(input));

        // calculate the eigenvalue decomposition
        EigenValueDecompositionSymm eigenDeco = new EigenValueDecompositionSymm(covarianceMatrix);
        E = eigenDeco.getEigenVectors();

        eigenValues = eigenDeco.getEigenValues();
        // calculate the resulting vectors
        resultVectors = Matrix.mult(Matrix.transpose(E), zeroMeanVectors);


        whiteningMatrix = Matrix.mult(Matrix.diag(Vector.invVector(Matrix.sqrtVector(eigenValues))), Matrix.transpose(E));
        dewhiteningMatrix = Matrix.mult(E, Matrix.diag(Matrix.sqrtVector(eigenValues)));

        // the whitened vectors' correlation matrix equals unit matrix
        // which is demanded to perform the FastICA algorithm
        whitenedVectors =
                Matrix.mult(whiteningMatrix, zeroMeanVectors);

    }

    /**
     * Calculates the mean vector from a set of vectors.
     * @param inVectors the set of vectors
     * @return the resulting mean vector
     */
    private static double[] calcMeanValues(double[][] vectors) {
        int m = Matrix.getNumOfRows(vectors);
        int n = Matrix.getNumOfColumns(vectors);
        double[] means = Vector.newVector(m);
        for (int i = 0; i < m; ++i) {
            means[i] = 0.0;
            for (int j = 0; j < n; ++j) {
                means[i] += vectors[i][j];
            }
            means[i] /= n;
        }
        return (means);
    }

    /**
     * Calculates a difference measure of two matrices
     * relative to their size.
     * @param mat1 the first matrix
     * @param mat2 the second matrix
     * @return the difference measure
     */
    private static double deltaMatrices(double[][] mat1, double[][] mat2) {
        double[][] test = Matrix.sub(mat1, mat2);
        double delta = 0.0;
        int m = Matrix.getNumOfRows(mat1);
        int n = Matrix.getNumOfColumns(mat1);
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                delta += Math.abs(test[i][j]);
            }
        }
        return (delta / (m * n));
    }

    /**
     * Calculates the power of a symmetric matrix.
     * @param inMatrix the symmetric matrix
     * @param power the power
     * @return the resulting matrix
     */
    private static double[][] powerSymmMatrix(double[][] inMatrix, double power) {
        EigenValueDecompositionSymm eigenDeco = new EigenValueDecompositionSymm(inMatrix);
        int m = Matrix.getNumOfRows(inMatrix);
        double[][] eigenVectors = eigenDeco.getEigenVectors();
        double[] eigenValues = eigenDeco.getEigenValues();
        for (int i = 0; i < m; ++i) {
            eigenValues[i] = Math.pow(eigenValues[i], power);
        }

        return (Matrix.mult(Matrix.mult(eigenVectors, Matrix.diag(eigenValues)), Matrix.transpose(eigenVectors)));
    }
}
